"""
批量读取图片
检测人脸坐标
按照人脸为中心，图片最短的边为正方形进行裁剪

"""


import os
from PIL import Image
from facenet_pytorch import MTCNN
import torch

# 初始化人脸检测模型
device = 'cuda' if torch.cuda.is_available() else 'cpu'
mtcnn = MTCNN(keep_all=False, device=device)

input_dir = r"D:\pyimg\Philippines-man-unsplash-105"
output_dir = r"D:\pyimg\Philippines-man-unsplash-105-crop"
os.makedirs(output_dir, exist_ok=True)

# 偏移比例（用于居中脖子）
VERTICAL_OFFSET_RATIO = 0.15  # 人脸中心向下偏移 15%


def smart_crop_avatar(image_path, save_path):
    try:
        img = Image.open(image_path).convert('RGB')
    except Exception as e:
        print(f"❌ 打开失败: {image_path} - {e}")
        return

    width, height = img.size
    crop_size = min(width, height)

    boxes, _ = mtcnn.detect(img)
    if boxes is not None:
        # 检测到人脸时
        x1, y1, x2, y2 = boxes[0]
        face_cx = (x1 + x2) / 2
        face_cy = (y1 + y2) / 2

        # 偏移以脖子为中心
        center_x = face_cx
        center_y = face_cy + crop_size * VERTICAL_OFFSET_RATIO
    else:
        # 未检测到人脸，使用图像中心裁剪
        print(f"⚠️ 未检测到人脸，居中裁剪: {os.path.basename(image_path)}")
        center_x = width / 2
        center_y = height / 2

    # 计算裁剪区域
    left = int(center_x - crop_size / 2)
    top = int(center_y - crop_size / 2)
    right = left + crop_size
    bottom = top + crop_size

    # 边界修正
    if left < 0:
        left = 0
        right = crop_size
    if top < 0:
        top = 0
        bottom = crop_size
    if right > width:
        right = width
        left = width - crop_size
    if bottom > height:
        bottom = height
        top = height - crop_size

    crop_box = (int(left), int(top), int(right), int(bottom))
    cropped_img = img.crop(crop_box)
    cropped_img.save(save_path)
    print(f"✅ 裁剪完成: {os.path.basename(save_path)}")

# 批量处理目录下所有图片
for filename in os.listdir(input_dir):
    if filename.lower().endswith(('.jpg', '.jpeg', '.png')):
        src_path = os.path.join(input_dir, filename)
        dst_path = os.path.join(output_dir, filename)
        smart_crop_avatar(src_path, dst_path)